"""
测试数据库工具模块
"""
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))

from api.controllers.utils.dmdb_utils import get_database_connection, create_sql_chain, execute_sql_query
from langchain_openai import ChatOpenAI

def test_database_connection():
    """测试数据库连接"""
    print("测试数据库连接...")
    db = get_database_connection()
    if db is None:
        print("数据库连接失败")
        return False
    else:
        print("数据库连接成功")
        # 显示可用的表
        try:
            tables = db.get_usable_table_names()
            print(f"可用的表: {tables}")
        except Exception as e:
            print(f"获取表信息失败: {e}")
        return True

def test_sql_chain():
    """测试SQL查询链"""
    print("\n测试SQL查询链...")
    
    # 创建模型
    try:
        chat_model = ChatOpenAI(
            model="gpt-3.5-turbo",
            temperature=0
        )
    except Exception as e:
        print(f"模型创建失败: {e}")
        return False
    
    # 创建SQL链
    chain = create_sql_chain(chat_model)
    if chain is None:
        print("SQL查询链创建失败")
        return False
    else:
        print("SQL查询链创建成功")
        return True

def test_sql_execution():
    """测试SQL执行"""
    print("\n测试SQL执行...")
    
    db = get_database_connection()
    if db is None:
        print("数据库连接失败")
        return False
    
    # 执行简单查询
    result = execute_sql_query("SELECT * FROM users LIMIT 1", db)
    if "error" in result:
        print(f"SQL执行失败: {result['error']}")
        return False
    else:
        print(f"SQL执行成功: {result['result']}")
        return True

def test_natural_language_query():
    """测试自然语言查询"""
    print("\n测试自然语言查询...")
    
    db = get_database_connection()
    if db is None:
        print("数据库连接失败")
        return False
    
    # 创建模型
    try:
        chat_model = ChatOpenAI(
            model="gpt-3.5-turbo",
            temperature=0
        )
    except Exception as e:
        print(f"模型创建失败: {e}")
        return False
    
    # 创建SQL链
    chain = create_sql_chain(chat_model)
    if chain is None:
        print("SQL查询链创建失败")
        return False
    
    # 使用自然语言查询
    question = "列出所有用户"
    try:
        sql_query = chain.invoke({"question": question})
        print(f"生成的SQL: {sql_query}")
        
        # 执行查询
        result = execute_sql_query(sql_query, db)
        if "error" in result:
            print(f"查询执行失败: {result['error']}")
            return False
        else:
            print(f"查询结果: {result['result']}")
            return True
    except Exception as e:
        print(f"自然语言查询失败: {e}")
        return False

if __name__ == "__main__":
    print("开始测试数据库工具模块...")
    
    # 测试数据库连接
    connection_success = test_database_connection()
    
    if connection_success:
        # 测试SQL链
        chain_success = test_sql_chain()
        
        # 测试SQL执行
        execution_success = test_sql_execution()
        
        # 测试自然语言查询
        nlq_success = test_natural_language_query()
        
        if chain_success and execution_success and nlq_success:
            print("\n所有测试通过!")
        else:
            print("\n部分测试失败!")
    else:
        print("数据库连接失败，跳过其他测试")